/*********************************************************************
RWLock: A simple and fast reader-writer lock class in C++
                   has characters of .NET ReaderWriterLock class
Copyright (C) 2006 Quynh Nguyen Huu

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
Lesser General Public License for more details.

Email questions, comments or suggestions to quynhnguyenhuu@gmail.com
*********************************************************************/

//#include "stdafx.h"
#include "rwlock.h"

#ifndef _ASSERT
#include <assert.h>
#define _ASSERT assert
#endif

HT_NAMESPACE_BEGIN

///////////////////////////////////////////////////////
// RWLock implementation

RWLockNoReentrance::RWLockNoReentrance()
: m_hSafeToWriteEvent(NULL), m_hSafeToReadEvent(NULL),
	m_iNumOfWriter(0), m_iNumOfReaderEntered(0), m_iNumOfReaderWaiting(0)
{
}

RWLockNoReentrance::~RWLockNoReentrance()
{
 	_ASSERT( (NULL == m_hSafeToReadEvent) &&
 		(NULL == m_hSafeToWriteEvent) );
}

bool RWLockNoReentrance::_ReaderWait(uint32 dwTimeout)
{
	bool blCanRead;

	++m_iNumOfReaderWaiting;
 	if(NULL == m_hSafeToReadEvent)
 	{
 		m_hSafeToReadEvent = new Event(false, true);	//CreateEvent(NULL, true, false, NULL);
 	}

	if(INFINITE == dwTimeout) // INFINITE is a special value
	{
		do
		{
			LeaveCS();
			//WaitForSingleObject(m_hSafeToReadEvent, INFINITE);
			m_hSafeToReadEvent->wait();
			// There might be one or more Writers entered, that's
			// why we need DO-WHILE loop here
			EnterCS();
		}
		while(0 != m_iNumOfWriter);

		++m_iNumOfReaderEntered;
		blCanRead = true;
	}
	else
	{
		LeaveCS();

		uint64 const dwBeginTime = HTGetTickCount();
		uint64 dwConsumedTime = 0;
		
		while(true)
		{
//			blCanRead = (WAIT_OBJECT_0 == WaitForSingleObject(m_hSafeToReadEvent,
//				dwTimeout - dwConsumedTime));
			blCanRead = m_hSafeToReadEvent->wait_with_timeout(dwTimeout - dwConsumedTime);

			EnterCS();

			if(0 == m_iNumOfWriter)
			{
				// Regardless timeout or not, there is no Writer
				// So it's safe to be Reader right now
				++m_iNumOfReaderEntered;
				blCanRead = true;
				break;
			}

			if(false == blCanRead)
			{
				// Timeout after waiting
				break;
			}

			// There are some Writers have just entered
			// So leave CS and prepare to try again
			LeaveCS();

			dwConsumedTime = HTGetTickCount() - dwBeginTime;
			if(dwConsumedTime > dwTimeout)
			{
				// Don't worry why the code here looks stupid
				// Because this case rarely happens, it's better
				//  to optimize code for the usual case
				blCanRead = false;
				EnterCS();
				break;
			}
		}
	}

	if(0==--m_iNumOfReaderWaiting)
	{
		// CloseHandle(m_hSafeToReadEvent);
		delete m_hSafeToReadEvent;
		m_hSafeToReadEvent = NULL;
	}

	return blCanRead;
}

void RWLockNoReentrance::_ReaderRelease()
{
	int32 _iNumOfReaderEntered = --m_iNumOfReaderEntered;
	_ASSERT(0 <= _iNumOfReaderEntered);

	if( (0 == _iNumOfReaderEntered) &&	(NULL != m_hSafeToWriteEvent) )
	{
		//SetEvent(m_hSafeToWriteEvent);
		m_hSafeToWriteEvent->signal();
	}
}

bool RWLockNoReentrance::_WriterWaitAndLeaveCSIfSuccess(uint32 dwTimeout)
{
	//EnterCS();
	_ASSERT(0 != dwTimeout);

	// Increase Writer-counter & reset Reader-event if necessary
	int32 _iNumOfWriter = ++m_iNumOfWriter;
	if(	(1 == _iNumOfWriter) && (NULL != m_hSafeToReadEvent) )
	{
		m_hSafeToReadEvent->signal();
	}

	if(NULL == m_hSafeToWriteEvent)
	{
		m_hSafeToWriteEvent = new Event(false, false);
	}

	LeaveCS();

	bool blCanWrite = m_hSafeToWriteEvent->wait_with_timeout(dwTimeout);
	if(false == blCanWrite)
	{
		// Undo what we changed after timeout
		EnterCS();
		if(0 == --m_iNumOfWriter)
		{
			//CloseHandle(m_hSafeToWriteEvent);
			delete m_hSafeToWriteEvent;
			m_hSafeToWriteEvent = NULL;

			if(0 == m_iNumOfReaderEntered)
			{
				// Although it was timeout, it's still safe to be writer now
				++m_iNumOfWriter;
				LeaveCS();
				blCanWrite = true;
			}
			else if(m_hSafeToReadEvent)
			{
				//SetEvent(m_hSafeToReadEvent);
				m_hSafeToReadEvent->signal();
			}
		}
	}
	return blCanWrite;
}

bool RWLockNoReentrance::_UpgradeToWriterLockAndLeaveCS(uint32 dwTimeout)
{
	_ASSERT(m_iNumOfReaderEntered > 0);

	if(0 == dwTimeout)
	{
		LeaveCS();
		return false;
	}

	--m_iNumOfReaderEntered;
	bool blCanWrite = _WriterWaitAndLeaveCSIfSuccess(dwTimeout);
	if(false == blCanWrite)
	{
		// Now analyze why it was failed to have suitable action
		if(0 == m_iNumOfWriter)
		{
			_ASSERT(0 < m_iNumOfReaderEntered);
			// There are some readers still owning the lock
			// It's safe to be a reader again after failure
			++m_iNumOfReaderEntered;
		}
		else
		{
			// Reach to here, it's NOT safe to be a reader immediately
			_ReaderWait(INFINITE);
			if(1 == m_iNumOfReaderEntered)
			{
				// After wait, now it's safe to be writer
				_ASSERT(0 == m_iNumOfWriter);
				m_iNumOfReaderEntered = 0;
				m_iNumOfWriter = 1;
				blCanWrite = true;
			}
		}
		LeaveCS();
	}

	return blCanWrite;
}

void RWLockNoReentrance::_WriterRelease(bool blDowngrade)
{
	_ASSERT(0 == m_iNumOfReaderEntered);

	if(blDowngrade)
	{
		++m_iNumOfReaderEntered;
	}

	if(0 == --m_iNumOfWriter)
	{

		if(NULL != m_hSafeToWriteEvent)
		{
			delete m_hSafeToWriteEvent;
			m_hSafeToWriteEvent = NULL;
		}

		if(m_hSafeToReadEvent)
		{
			m_hSafeToReadEvent->signal();
		}
	}
	else
	{
		//////////////////////////////////////////////////////////////////////////
		// Some WRITERs are queued
		_ASSERT( (0 < m_iNumOfWriter) && (NULL != m_hSafeToWriteEvent));

		if(false == blDowngrade)
		{
			m_hSafeToWriteEvent->signal();
		}
	}
}

bool RWLockNoReentrance::AcquireReaderLock(uint32 dwTimeout)
{
	bool blCanRead;

	EnterCS();
	if(0 == m_iNumOfWriter)
	{
		// Enter successful without wait
		++m_iNumOfReaderEntered;
		blCanRead = true;
	}
	else
	{
		 blCanRead = (dwTimeout)? _ReaderWait(dwTimeout) : false;
	}
	LeaveCS();

	return blCanRead;
}

void RWLockNoReentrance::ReleaseReaderLock()
{
	EnterCS();
	_ReaderRelease();
	LeaveCS();
}

bool RWLockNoReentrance::AcquireWriterLock(uint32 dwTimeout)
{
	bool blCanWrite ;

	EnterCS();
	if(0 == (m_iNumOfWriter | m_iNumOfReaderEntered))
	{
		++m_iNumOfWriter;
		blCanWrite = true;
	}
	else if(0 == dwTimeout)
	{
		blCanWrite = false;
	}
	else
	{
		blCanWrite = _WriterWaitAndLeaveCSIfSuccess(dwTimeout);
		if(blCanWrite)
		{
			return true;
		}
	}

	LeaveCS();
	return blCanWrite;
}

void RWLockNoReentrance::ReleaseWriterLock()
{
	EnterCS();
	_WriterRelease(false);
	LeaveCS();
}

void RWLockNoReentrance::DowngradeFromWriterLock()
{
	EnterCS();
	_WriterRelease(true);
	LeaveCS();
}

bool RWLockNoReentrance::UpgradeToWriterLock(uint32 dwTimeout)
{
	EnterCS();
	return _UpgradeToWriterLockAndLeaveCS(dwTimeout);
}

// END RWLockNoReentrance implementation
///////////////////////////////////////////////////////

///////////////////////////////////////////////////////
// RWLock implementation

#define READER_RECURRENCE_UNIT 0x00000001
#define READER_RECURRENCE_MASK 0x0000FFFF
#define WRITER_RECURRENCE_UNIT 0x00010000

RWLock::RWLock()
{
}

RWLock::~RWLock()
{
}

// ��������ʱ���������ӱ��̵߳ļ���ֵ
bool RWLock::AcquireReaderLock(uint32 dwTimeout)
{
	const uint32 dwCurrentThreadId = GetCurrentThreadId();

	EnterCS();
	CMapThreadToState::iterator ite = m_map.find(dwCurrentThreadId);

	if(ite != m_map.end())
	{
		//////////////////////////////////////////////////////////////////////////
		// Current thread was already a WRITER or READER
		_ASSERT(0 < ite->second);
		LeaveCS();

		ite->second += READER_RECURRENCE_UNIT;
		return true;
	}

	if(0 == m_iNumOfWriter)
	{
		// There is NO WRITER on this RW object
		// Current thread is going to be a READER
		++m_iNumOfReaderEntered;
		m_map.insert(std::make_pair(dwCurrentThreadId, READER_RECURRENCE_UNIT));

		LeaveCS();
		return true;
	}

	if(0 == dwTimeout)
	{
		LeaveCS();
		return false;
	}

	bool blCanRead = _ReaderWait(dwTimeout);
	if(blCanRead)
	{
		m_map.insert(std::make_pair(dwCurrentThreadId, READER_RECURRENCE_UNIT));
	}
	LeaveCS();

	return blCanRead;
}

// �ͷŶ���
void RWLock::ReleaseReaderLock()
{
	const uint32 dwCurrentThreadId = GetCurrentThreadId();
	EnterCS();

	CMapThreadToState::iterator ite = m_map.find(dwCurrentThreadId);
	_ASSERT( (ite != m_map.end()) && (READER_RECURRENCE_MASK & ite->second));

	const uint32 dwThreadState = (ite->second -= READER_RECURRENCE_UNIT);
	if(0 == dwThreadState)
	{
		m_map.erase(ite);
		_ReaderRelease();
	}
	LeaveCS();
}

// ����д��
bool RWLock::AcquireWriterLock(uint32 dwTimeout)
{
	const uint32 dwCurrentThreadId = GetCurrentThreadId();
	bool blCanWrite;

	EnterCS();
	CMapThreadToState::iterator ite = m_map.find(dwCurrentThreadId);

	if(ite != m_map.end())
	{// 
		_ASSERT(0 < ite->second);

		if(ite->second >= WRITER_RECURRENCE_UNIT)
		{// �Ѿ��ڱ��߳���д״̬��,����д����
			// Current thread was already a WRITER
			LeaveCS();
			ite->second += WRITER_RECURRENCE_UNIT;
			return true;
		}

		// Current thread was already a READER
		_ASSERT(1 <= m_iNumOfReaderEntered);
		if(1 == m_iNumOfReaderEntered)
		{// 
			// This object is owned by ONLY current thread for READ
			// There might be some threads queued to be WRITERs
			// but for effectiveness (higher throughput), we allow current
			// thread upgrading to be WRITER right now
			m_iNumOfReaderEntered = 0;
			++m_iNumOfWriter;
			LeaveCS();
			ite->second += WRITER_RECURRENCE_UNIT;
			return true;
		}

		// Try upgrading from reader to writer
		blCanWrite = _UpgradeToWriterLockAndLeaveCS(dwTimeout);
		if(blCanWrite)
		{
			ite->second += WRITER_RECURRENCE_UNIT;
		}
	}
	else
	{
		if(0 == (m_iNumOfWriter | m_iNumOfReaderEntered))
		{
			// This RW object is not owned by any thread
			// --> it's safe to make this thread to be WRITER
			++m_iNumOfWriter;
			m_map.insert(std::make_pair(dwCurrentThreadId, WRITER_RECURRENCE_UNIT));
			LeaveCS();
			return true;
		}

		if(0 == dwTimeout)
		{
			LeaveCS();
			return false;
		}

		blCanWrite = _WriterWaitAndLeaveCSIfSuccess(dwTimeout);
		if(blCanWrite)
		{
			EnterCS();
			m_map.insert(std::make_pair(dwCurrentThreadId, WRITER_RECURRENCE_UNIT));
		}
		LeaveCS();
	}

	return blCanWrite;
}

void RWLock::ReleaseWriterLock()
{
	const uint32 dwCurrentThreadId = GetCurrentThreadId();
	EnterCS();

	CMapThreadToState::iterator ite = m_map.find(dwCurrentThreadId);
	_ASSERT( (ite != m_map.end()) && (WRITER_RECURRENCE_UNIT <= ite->second));

	const uint32 dwThreadState = (ite->second -= WRITER_RECURRENCE_UNIT);
	if(0 == dwThreadState)
	{
		m_map.erase(ite);
		_WriterRelease(false);
	}
	else if (WRITER_RECURRENCE_UNIT > dwThreadState)
	{
		// Down-grading from writer to reader
		_WriterRelease(true);
	}
	LeaveCS();
}

void RWLock::ReleaseAllLocks()
{
	const uint32 dwCurrentThreadId = GetCurrentThreadId();

	EnterCS();
	CMapThreadToState::iterator ite = m_map.find(dwCurrentThreadId);
	if(ite != m_map.end())
	{
		const uint32 dwThreadState = ite->second;
		m_map.erase(ite);
		if(WRITER_RECURRENCE_UNIT <= dwThreadState)
		{
			_WriterRelease(false);
		}
		else
		{
			_ASSERT(0 < dwThreadState);
			_ReaderRelease();
		}
	}
	LeaveCS();
}

uint32 RWLock::GetCurrentThreadStatus() const
{
	uint32 dwThreadState;
	const uint32 dwCurrentThreadId = GetCurrentThreadId();

	EnterCS();
	CMapThreadToState::const_iterator ite = m_map.find(dwCurrentThreadId);
	if(ite != m_map.end())
	{
		LeaveCS();
		dwThreadState = ite->second;
		_ASSERT(dwThreadState > 0);
	}
	else
	{
		LeaveCS();
		dwThreadState = 0;
	}

	return dwThreadState;
}

void RWLock::GetCurrentThreadStatus(uint32* lpdwReaderLockCounter,
											   uint32* lpdwWriterLockCounter) const
{
	const uint32 dwThreadState = GetCurrentThreadStatus();

	if(NULL != lpdwReaderLockCounter)
	{
		*lpdwReaderLockCounter = (dwThreadState & READER_RECURRENCE_MASK);
	}

	if(NULL != lpdwWriterLockCounter)
	{
		*lpdwWriterLockCounter = (dwThreadState / WRITER_RECURRENCE_UNIT);
	}
}

HT_NAMESPACE_END

